#ifndef FLOW_H
#define FLOW_H

#include <cstdio>
#include <vector>
#include <cutil_inline.h>
#include <vector_types.h>

#ifndef CUDA_SAFE_FREE
#define CUDA_SAFE_FREE(ptr) { if(ptr) { CUDA_SAFE_CALL(cudaFree(ptr)); ptr = NULL; } }
#endif

template<typename T>
extern void initDensity(int nx, int ny, int nz, T *d);

template<typename T>
void initVelocity(int nx, int ny, int nz, T *u, T *v, T *w);

template<typename T>
void addSource(int nx, int ny, int nz, T dx, T dy, T dz, T dt, T *u, T *v, T *w, T *d, T *t);

template<typename T>
void convect(int nx, int ny, int nz, T dx, T dy, T dz, T dt, T *u, T *v, T *w, T *f, T *fout);

template<typename T>
void diffuse(int nx, int ny, int nz, T dx, T dy, T dz, T dt, T nu, T *f, T *fout, int its);

template<typename T>
void divergence(int nx, int ny, int nz, T dx, T dy, T dz, T *u, T *v, T *w, T *fout);

template<typename T>
void poisson(int nx, int ny, int nz, T dx, T dy, T dz, T *d, T *p, int its);

template<typename T>
void project(int nx, int ny, int nz, T dx, T dy, T dz, T *u, T *v, T *w, T *p);

template<typename T>
void setBoundary(int nx, int ny, int nz, T dx, T dy, T dz, T *f, T value);

template<typename T>
void setNaturalBoundary(int nx, int ny, int nz, T dx, T dy, T dz, T *f);

template<typename T>
void setReflectBoundary(int nx, int ny, int nz, T dx, T dy, T dz, T *u, T *v, T *w);

template<typename T>
struct FlowConfig
{
public:
	// grid size
	int nx;
	int ny;
	int nz;
	// cell size
	T dx;
	T dy;
	T dz;
	// physical parameters
	T dt;	// time step
	T rho;	// density
	T mu;	// viscosity
	T gx;	// gravity x
	T gy;	// gravity y
	T gz;	// gravity z
	T kappa;	// thermal diffusivity

public:
	FlowConfig()
	{
		nx = 9;
		ny = 9;
		nz = 9;
		dx = T(1) / T(nx - 1);
		dy = T(1) / T(ny - 1);
		dz = T(1) / T(nz - 1);
		dt = T(0.001);
		rho = T(1);
		mu = T(0.001);
		gx = T(0);
		gy = T(0);
		gz = T(0);
		kappa = T(0.00001);
	}
};

template<typename T>
struct FlowSolution
{
public:
	T *u;
	T *v;
	T *w;
	T *d;
	T *t;

public:
	FlowSolution()
	{
		u = NULL;
		v = NULL;
		w = NULL;
		d = NULL;
		t = NULL;
	}

	~FlowSolution()
	{
		free();
	}

public:
	void free()
	{
		CUDA_SAFE_FREE(u);
		CUDA_SAFE_FREE(v);
		CUDA_SAFE_FREE(w);
		CUDA_SAFE_FREE(d);
		CUDA_SAFE_FREE(t);
	}

	void resize(int size)
	{
		free();
		CUDA_SAFE_CALL(cudaMalloc((void **)&u, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMalloc((void **)&v, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMalloc((void **)&w, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMalloc((void **)&d, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMalloc((void **)&t, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(u, 0, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(v, 0, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(w, 0, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(d, 0, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(t, 0, size * sizeof(T)));
	}
};

template<typename T>
struct FlowProjection
{
public:
	T *d;	// divergence
	T *p;	// pressure

public:
	FlowProjection()
	{
		d = NULL;
		p = NULL;
	}

	~FlowProjection()
	{
		free();
	}

public:
	void free()
	{
		CUDA_SAFE_FREE(d);
		CUDA_SAFE_FREE(p);
	}

	void resize(int size)
	{
		free();
		CUDA_SAFE_CALL(cudaMalloc((void **)&d, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMalloc((void **)&p, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(d, 0, size * sizeof(T)));
		CUDA_SAFE_CALL(cudaMemset(p, 0, size * sizeof(T)));
	}

	void apply(const FlowConfig<T> &c, FlowSolution<T> &s, int its)
	{
		int size = c.nx * c.ny * c.nz;
		divergence(c.nx, c.ny, c.nz, c.dx, c.dy, c.dz, s.u, s.v, s.w, d);
		poisson(c.nx, c.ny, c.nz, c.dx, c.dy, c.dz, d, p, its);
		project(c.nx, c.ny, c.nz, c.dx, c.dy, c.dz, s.u, s.v, s.w, p);
		setReflectBoundary(c.nx, c.ny, c.nz, c.dx, c.dy, c.dz, s.u, s.v, s.w);
	}
};

template<typename T>
class Flow
{
protected:
	int steps;
	T time;
	FlowConfig<T> config;
	FlowSolution<T> solution[2];
	FlowProjection<T> projection;

protected:
	void debug(const char *name, T *data)
	{
		int size = config.nx * config.ny * config.nz;
		std::vector<T> tmp(size);
		cudaMemcpy(&tmp.front(), data, size * sizeof(T), cudaMemcpyDeviceToHost);
		puts(name);
		for(int i = 0; i < size; ++i)
		{
			printf("%+.4f ", tmp[i]);
			if((i + 1) % (config.nx) == 0)
				putchar('\n');
			if((i + 1) % (config.nx * config.ny) == 0)
				putchar('\n');
		}
	}

public:
	void dump(std::vector<T> &u, std::vector<T> &v, std::vector<T> &w, std::vector<T> &d)
	{
		int size = config.nx * config.ny * config.nz;
		u.resize(size);
		v.resize(size);
		w.resize(size);
		d.resize(size);
		cudaMemcpy(&u.front(), solution[0].u, size * sizeof(float), cudaMemcpyDeviceToHost);
		cudaMemcpy(&v.front(), solution[0].v, size * sizeof(float), cudaMemcpyDeviceToHost);
		cudaMemcpy(&w.front(), solution[0].w, size * sizeof(float), cudaMemcpyDeviceToHost);
		cudaMemcpy(&d.front(), solution[0].d, size * sizeof(float), cudaMemcpyDeviceToHost);
	}

	void dump2(std::vector<T> &d, std::vector<T> &p)
	{
		int size = config.nx * config.ny * config.nz;
		d.resize(size);
		p.resize(size);
		cudaMemcpy(&d.front(), projection.d, size * sizeof(float), cudaMemcpyDeviceToHost);
		cudaMemcpy(&p.front(), projection.p, size * sizeof(float), cudaMemcpyDeviceToHost);
	}

	void init(const FlowConfig<T> &flowConfig)
	{
		steps = 0;
		time = T(0);
		config = flowConfig;
		int size = config.nx * config.ny * config.nz;
		solution[0].resize(size);
		solution[1].resize(size);
		projection.resize(size);

		//initDensity(config.nx, config.ny, config.nz, solution[0].d);
		//debug("density0:", solution[steps & 1].d);

		//initVelocity(config.nx, config.ny, config.nz, solution[0].u, solution[0].v, solution[0].w);
		//debug("velocityU:", solution[steps & 1].u);
	}

	void step()
	{
		const int its = 32;
		const int idx0 = 0;
		const int idx1 = 1;
		T nu = config.mu / config.rho;

		// convect velocity
		convect(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].u, solution[idx1].u);
		convect(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].v, solution[idx1].v);
		convect(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].w, solution[idx1].w);
		setReflectBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx1].u, solution[idx1].v, solution[idx1].w);

		// diffuse velocity
		diffuse(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, nu, solution[idx1].u, solution[idx0].u, its);
		diffuse(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, nu, solution[idx1].v, solution[idx0].v, its);
		diffuse(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, nu, solution[idx1].w, solution[idx0].w, its);
		setReflectBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx0].u, solution[idx0].v, solution[idx0].w);

		// source
		addSource(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].d, solution[idx0].t);

		// project velocity
		projection.apply(config, solution[idx0], its);

		// convect density & temperature
		convect(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].d, solution[idx1].d);
		convect(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, solution[idx0].u, solution[idx0].v, solution[idx0].w, solution[idx0].t, solution[idx1].t);
		setNaturalBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx1].d);
		setNaturalBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx1].t);

		// diffuse density & temperature
		diffuse(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, T(0), solution[idx1].d, solution[idx0].d, its);
		diffuse(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, config.dt, config.kappa, solution[idx1].t, solution[idx0].t, its);
		setNaturalBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx0].d);
		setNaturalBoundary(config.nx, config.ny, config.nz, config.dx, config.dy, config.dz, solution[idx0].t);

		++steps;
		time += config.dt;
		//debug("density:", solution[idx0].d);
	}
};

#endif
